#include "hero_integrated_vision/number_classifier.hpp"

#include <openvino/openvino.hpp>
#include <opencv2/opencv.hpp>

#include <rclcpp/rclcpp.hpp>
#include <algorithm>
#include <fstream>
#include <vector>

namespace rm_aim
{

NumberClassifier::NumberClassifier(
  const std::string & model_path,
  const std::string & label_path,
  const double threshold,
  const rclcpp::Logger & logger,
  const std::vector<std::string> & ignore_classes)
: threshold_(threshold),
  logger_(logger),
  ignore_classes_(ignore_classes)
{
  try {
    // 初始化OpenVINO核心
    core_ = std::make_unique<ov::Core>();
    
    // 读取模型
    RCLCPP_INFO(logger_, "Loading model: %s", model_path.c_str());
    model_ = core_->read_model(model_path);
    
    // 检查输入输出数量
    if (model_->inputs().size() != 1) {
      throw std::invalid_argument("Invalid model input count");
    }
    if (model_->outputs().size() != 1) {
      throw std::invalid_argument("Invalid model output count");
    }

    // 读取标签文件
    loadLabels(label_path);
    
    RCLCPP_INFO(logger_, "Number classifier initialized successfully");
  } catch (const std::exception & e) {
    RCLCPP_FATAL(logger_, "Initialization failed: %s", e.what());
    throw;
  }
}

void NumberClassifier::loadLabels(const std::string & label_path)
{
  std::ifstream label_file(label_path);
  if (!label_file.is_open()) {
    throw std::runtime_error("Failed to open label file");
  }

  std::string line;
  while (std::getline(label_file, line)) {
    class_names_.push_back(line);
  }
  
  RCLCPP_INFO(logger_, "Number of loaded labels: %zu", class_names_.size());
}

void NumberClassifier::compileModel(const std::string & device)
{
  try {
    RCLCPP_INFO(logger_, "Compiling model on device: %s", device.c_str());
    compiled_model_ = core_->compile_model(model_, device);
    infer_request_ = compiled_model_.create_infer_request();
    RCLCPP_INFO(logger_, "Model compiled successfully");
  } catch (const std::exception & e) {
    RCLCPP_ERROR(logger_, "Model compilation failed: %s", e.what());
    throw;
  }
}

void NumberClassifier::extractNumbers(const cv::Mat & src, std::vector<Armor> & armors)
{
  // 数字区域提取逻辑
  const int light_length = 12;
  const int warp_height = 28;
  const int small_armor_width = 32;
  const int large_armor_width = 54;
  const cv::Size roi_size(20, 28);

  for (auto & armor : armors) {
    cv::Point2f lights_vertices[4] = {
      armor.left_light.bottom_point, armor.left_light.top_point, 
      armor.right_light.top_point, armor.right_light.bottom_point
    };

    const int top_light_y = (warp_height - light_length) / 2 - 1;
    const int bottom_light_y = top_light_y + light_length;
    const int warp_width = armor.type == ArmorType::SMALL ? 
                          small_armor_width : large_armor_width;
                          
    cv::Point2f target_vertices[4] = {
      cv::Point(0, bottom_light_y),
      cv::Point(0, top_light_y),
      cv::Point(warp_width - 1, top_light_y),
      cv::Point(warp_width - 1, bottom_light_y),
    };

    cv::Mat number_image;
    auto rotation_matrix = cv::getPerspectiveTransform(lights_vertices, target_vertices);
    cv::warpPerspective(src, number_image, rotation_matrix, cv::Size(warp_width, warp_height));

    // 提取ROI区域
    number_image = number_image(cv::Rect(
      cv::Point((warp_width - roi_size.width) / 2, 0), 
      roi_size
    ));

    // 二值化处理
    cv::cvtColor(number_image, number_image, cv::COLOR_RGB2GRAY);
    cv::threshold(number_image, number_image, 0, 255, 
                 cv::THRESH_BINARY | cv::THRESH_OTSU);

    armor.number_roi_img = number_image;
  }
}

void NumberClassifier::classify(std::vector<Armor> & armors)
{
  if (armors.empty()) return;

  try {
    for (auto & armor : armors) {
      if (armor.number_roi_img.empty()) {
        RCLCPP_WARN(logger_, "Armor number image is empty, skipping classification");
        continue;
      }

      // 准备输入数据
      auto input_tensor = prepareInputTensor(armor.number_roi_img);
      
      // 设置输入并执行推理
      infer_request_.set_input_tensor(input_tensor);
      infer_request_.infer();
      
      // 处理输出结果
      processOutput(armor);
    }

    // 过滤无效装甲板
    filterArmors(armors);
  } catch (const std::exception & e) {
    RCLCPP_ERROR(logger_, "Classification process error: %s", e.what());
  }
}

ov::Tensor NumberClassifier::prepareInputTensor(const cv::Mat & image)
{
  // 归一化处理
  cv::Mat normalized;
  image.convertTo(normalized, CV_32F, 1.0 / 255.0);

  // 设置输入形状和布局
  const ov::Shape input_shape = {1, 1, static_cast<size_t>(normalized.rows), static_cast<size_t>(normalized.cols)};
  const ov::element::Type input_type = ov::element::f32;

  // 创建输入张量
  ov::Tensor input_tensor(input_type, input_shape);
  float * input_data = input_tensor.data<float>();

  // 复制数据到张量 (NCHW格式)
  for (int i = 0; i < normalized.rows; ++i) {
    for (int j = 0; j < normalized.cols; ++j) {
      input_data[i * normalized.cols + j] = normalized.at<float>(i, j);
    }
  }

  return input_tensor;
}

void NumberClassifier::processOutput(Armor & armor)
{  // 获取输出张量
  const ov::Tensor & output_tensor = infer_request_.get_output_tensor();
  // 修改为使用缓冲区方式来获取输出数据
  const ov::Shape output_shape = output_tensor.get_shape();
  const size_t output_size = output_tensor.get_size();

  // 创建缓冲区并复制输出数据
  std::vector<float> output_buffer(output_size); // 缓冲区（大小=输出元素数）
  ov::Tensor buffer_tensor(ov::element::f32, output_shape, output_buffer.data()); // 绑定缓冲区的张量
  output_tensor.copy_to(buffer_tensor); // 复制数据到缓冲区
  const float * output_data = output_buffer.data();

  // 计算softmax
  std::vector<float> probabilities(output_size);
  computeSoftmax(output_data, output_size, probabilities.data());

  // 找到最大概率的类别
  auto max_it = std::max_element(probabilities.begin(), probabilities.end());
  armor.classification_conf = *max_it;
  const int class_idx = std::distance(probabilities.begin(), max_it);

  // 检查类别索引有效性
  if (class_idx >= 0 && static_cast<size_t>(class_idx) < class_names_.size()) {
    armor.number = class_names_[class_idx];
  } else {
    armor.number = "unknown";
    RCLCPP_WARN(logger_, "Invalid class index: %d", class_idx);
  }

  // 生成分类结果字符串
  std::stringstream ss;
  ss << (armor.left_light.camp_color == RED_CAMP ? "red" : "blue") 
   << ", " 
   << armor.number 
   << ": " 
   << std::fixed << std::setprecision(1) 
   << armor.classification_conf * 100.0 
   << "%";
armor.classification_result = ss.str();
}

void NumberClassifier::computeSoftmax(const float * input, size_t size, float * output)
{
  if (size == 0) return;

  // 数值稳定化：减去最大值
  const float max_val = *std::max_element(input, input + size);
  
  // 计算指数
  float sum = 0.0f;
  for (size_t i = 0; i < size; ++i) {
    output[i] = std::exp(input[i] - max_val);
    sum += output[i];
  }
  
  // 归一化
  for (size_t i = 0; i < size; ++i) {
    output[i] /= sum;
  }
}

void NumberClassifier::filterArmors(std::vector<Armor> & armors)
{
  armors.erase(
    std::remove_if(
      armors.begin(), armors.end(),
      [this](const Armor & armor) {
        // 过滤低置信度结果
        if (armor.classification_conf < threshold_) {
          RCLCPP_DEBUG(logger_, "Filtering low-confidence armor: %s", armor.classification_result.c_str());
          return true;
        }

        // 过滤忽略类别
        if (std::find(ignore_classes_.begin(), ignore_classes_.end(), armor.number) != 
            ignore_classes_.end()) {
          RCLCPP_DEBUG(logger_, "Filtering ignored class: %s", armor.number.c_str());
          return true;
        }

        // 过滤类型不匹配的装甲板
        bool type_mismatch = false;
        if (armor.type == ArmorType::LARGE) {
          type_mismatch = (armor.number == "outpost" || armor.number == "2" || 
                         armor.number == "guard");
        } else if (armor.type == ArmorType::SMALL) {
          type_mismatch = (armor.number == "1" || armor.number == "base");
        }

        if (type_mismatch) {
          RCLCPP_DEBUG(logger_, "Filtering type-mismatched armor: %s", armor.number.c_str());
          return true;
        }

        return false;
      }),
    armors.end());
}

}  // namespace rm_aim